{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KubeFlow Pipeline: Github Issue Summarization using Tensor2Tensor\n",
    "\n",
    "This notebook assumes that you have already set up a GKE cluster with Kubeflow installed as per this codelab: [g.co/codelabs/kubecon18](g.co/codelabs/kubecon18). Currently, this notebook must be run from the Kubeflow JupyterHub installation, as described in the codelab.\n",
    "\n",
    "In this notebook, we will show how to:\n",
    "\n",
    "* Interactively define a KubeFlow Pipeline using the Pipelines Python SDK\n",
    "* Submit and run the pipeline\n",
    "* Add a step in the pipeline\n",
    "\n",
    "This example pipeline trains a [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor/) model on Github issue data, learning to predict issue titles from issue bodies. It then exports the trained model and deploys the exported model to [Tensorflow Serving](https://github.com/tensorflow/serving). \n",
    "The final step in the pipeline launches a web app which interacts with the TF-Serving instance in order to get model predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do some imports and set some variables.  Set the `WORKING_DIR` to a path under the Cloud Storage bucket you created earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp  # the Pipelines SDK.  This library is included with the notebook image.\n",
    "from kfp import compiler\n",
    "import kfp.dsl as dsl\n",
    "import kfp.gcp as gcp\n",
    "import kfp.notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define some pipeline input variables. \n",
    "WORKING_DIR = 'gs://YOUR_GCS_BUCKET/t2t/notebooks' # Such as gs://bucket/object/path\n",
    "PROJECT_NAME = 'YOUR_PROJECT'\n",
    "GITHUB_TOKEN = 'YOUR_GITHUB_TOKEN'  # needed for prediction, to grab issue data from GH\n",
    "\n",
    "DEPLOY_WEBAPP = 'false'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an *Experiment* in the Kubeflow Pipeline System\n",
    "\n",
    "The Kubeflow Pipeline system requires an \"Experiment\" to group pipeline runs. You can create a new experiment, or call `client.list_experiments()` to get existing ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that this notebook should be running in JupyterHub in the same cluster as the pipeline system.\n",
    "# Otherwise, additional config would be required to connect.\n",
    "client = kfp.Client()\n",
    "client.list_experiments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp = client.create_experiment(name='datagen_notebook')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a Pipeline\n",
    "\n",
    "Authoring a pipeline is like authoring a normal Python function. The pipeline function describes the topology of the pipeline. \n",
    "\n",
    "Each step in the pipeline is typically a `ContainerOp` --- a simple class or function describing how to interact with a docker container image. In the pipeline, all the container images referenced in the pipeline are already built. \n",
    "\n",
    "The pipeline starts by training a [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor/) model, using already-preprocessed data. (More accurately, this step starts from an existing model checkpoint, then trains for a few more hundred steps).  When it finishes, it exports the model in a form suitable for serving by [TensorFlow serving](https://github.com/tensorflow/serving/).\n",
    "\n",
    "The next step deploys a TF-serving instance with that model.\n",
    "\n",
    "The last step launches a web app with which you can interact with the TF-serving instance to get model predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@dsl.pipeline(\n",
    "  name='Github issue summarization',\n",
    "  description='Demonstrate Tensor2Tensor-based training and TF-Serving'\n",
    ")\n",
    "def gh_summ(\n",
    "  train_steps = 2019300,\n",
    "  project = PROJECT_NAME,\n",
    "  github_token = GITHUB_TOKEN,\n",
    "  working_dir = GITHUB_TOKEN,\n",
    "  checkpoint_dir = 'gs://aju-dev-demos-codelabs/kubecon/model_output_tbase.bak2019000',\n",
    "  deploy_webapp = 'true',\n",
    "  data_dir = 'gs://aju-dev-demos-codelabs/kubecon/t2t_data_gh_all/'\n",
    "  ):\n",
    "\n",
    "  train = dsl.ContainerOp(\n",
    "      name = 'train',\n",
    "      image = 'gcr.io/google-samples/ml-pipeline-t2ttrain:v1ap',\n",
    "      arguments = [ \"--data-dir\", data_dir,\n",
    "          \"--checkpoint-dir\", checkpoint_dir,\n",
    "          \"--working-dir\", working_dir,\n",
    "          \"--model-dir\", '%s/%s/model_output' % (working_dir, '{{workflow.name}}'),\n",
    "          \"--train-steps\", train_steps, \"--deploy-webapp\" , deploy_webapp],\n",
    "      file_outputs={'output': '/tmp/output'}\n",
    "\n",
    "      ).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "\n",
    "  serve = dsl.ContainerOp(\n",
    "      name = 'serve',\n",
    "      image = 'gcr.io/google-samples/ml-pipeline-kubeflow-tfserve',\n",
    "      arguments = [\"--model_name\", 'ghsumm-%s' % ('{{workflow.name}}',),\n",
    "          \"--model_path\", '%s/%s/model_output/export' % (working_dir, '{{workflow.name}}')\n",
    "          ]\n",
    "      )\n",
    "  serve.after(train)\n",
    "  train.set_gpu_limit(4)\n",
    "  train.set_memory_limit('48G')\n",
    "\n",
    "  with dsl.Condition(train.output=='true'):\n",
    "    webapp = dsl.ContainerOp(\n",
    "        name = 'webapp',\n",
    "        image = 'gcr.io/google-samples/ml-pipeline-webapp-launcher',\n",
    "        arguments = [\"--model_name\", 'ghsumm-%s' % ('{{workflow.name}}',),\n",
    "            \"--github_token\", github_token]\n",
    "\n",
    "        )\n",
    "    webapp.after(serve)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Submit an experiment *run*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiler.Compiler().compile(gh_summ, 'ghsumm.tar.gz')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The call below will run the compiled pipeline.  We won't actually do that now, but instead we'll add a new step to the pipeline, then run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You'd uncomment this call to actually run the pipeline. \n",
    "# run = client.run_pipeline(exp.id, 'ghsumm', 'ghsumm.tar.gz',\n",
    "#                           params={'working-dir': WORKING_DIR,\n",
    "#                                   'github-token': GITHUB_TOKEN,\n",
    "#                                   'project': PROJECT_NAME})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add a step to the pipeline\n",
    "\n",
    "Next, let's add a new step to the pipeline above.  As currently defined, the pipeline accesses a directory of pre-processed data as input to training.  Let's see how we could include the pre-processing as part of the pipeline. \n",
    "\n",
    "We're going to cheat a bit, as processing the full dataset will take too long for this workshop, so we'll use a smaller sample. For that reason, you won't actually make use of the generated data from this step (we'll stick to using the full dataset for training), but this shows how you could do so if we had more time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll define the new pipeline step. Note the last line of this new function, which gives this step's pod the credentials to access GCP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining the new data preprocessing pipeline step. \n",
    "# Note the last line, which gives this step's pod the credentials to access GCP\n",
    "def preproc_op(data_dir, project):\n",
    "  return dsl.ContainerOp(\n",
    "    name='datagen',\n",
    "    image='gcr.io/google-samples/ml-pipeline-t2tproc',\n",
    "    arguments=[ \"--data-dir\", data_dir, \"--project\", project]\n",
    "  ).apply(gcp.use_gcp_secret('user-gcp-sa'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modify the pipeline to add the new step\n",
    "\n",
    "Now, we'll redefine the pipeline to add the new step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Then define a new Pipeline. It's almost the same as the original one, \n",
    "# but with the addition of the data processing step.\n",
    "import kfp.dsl as dsl\n",
    "import kfp.gcp as gcp\n",
    "\n",
    "@dsl.pipeline(\n",
    "  name='Github issue summarization',\n",
    "  description='Demonstrate TFT-based feature processing, TFMA, TFJob, CMLE OP, and TF-Serving'\n",
    ")\n",
    "def gh_summ2(\n",
    "  train_steps = 2019300,\n",
    "  project = PROJECT_NAME,\n",
    "  github_token = GITHUB_TOKEN,\n",
    "  working_dir = GITHUB_TOKEN,\n",
    "  checkpoint_dir = 'gs://aju-dev-demos-codelabs/kubecon/model_output_tbase.bak2019000',\n",
    "  deploy_webapp = 'true',\n",
    "  data_dir = 'gs://aju-dev-demos-codelabs/kubecon/t2t_data_gh_all/'\n",
    "  ):\n",
    "\n",
    "  # The new pre-processing op.\n",
    "  preproc = preproc_op(project=project,\n",
    "      data_dir=('%s/%s/gh_data' % (working_dir, '{{workflow.name}}')))\n",
    "\n",
    "  train = dsl.ContainerOp(\n",
    "      name = 'train',\n",
    "      image = 'gcr.io/google-samples/ml-pipeline-t2ttrain:v1ap',\n",
    "      arguments = [ \"--data-dir\", data_dir,\n",
    "          \"--checkpoint-dir\", checkpoint_dir,\n",
    "          \"--working-dir\", working_dir,\n",
    "          \"--model-dir\", '%s/%s/model_output' % (working_dir, '{{workflow.name}}'),\n",
    "          \"--train-steps\", train_steps, \"--deploy-webapp\" , deploy_webapp],\n",
    "      file_outputs={'output': '/tmp/output'}\n",
    "\n",
    "      ).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "  train.after(preproc)\n",
    "\n",
    "  serve = dsl.ContainerOp(\n",
    "      name = 'serve',\n",
    "      image = 'gcr.io/google-samples/ml-pipeline-kubeflow-tfserve',\n",
    "      arguments = [\"--model_name\", 'ghsumm-%s' % ('{{workflow.name}}',),\n",
    "          \"--model_path\", '%s/%s/model_output/export' % (working_dir, '{{workflow.name}}')\n",
    "          ]\n",
    "      )\n",
    "  serve.after(train)\n",
    "  train.set_gpu_limit(4)\n",
    "  train.set_memory_limit('48G')    \n",
    "\n",
    "  with dsl.Condition(train.output=='true'):\n",
    "    webapp = dsl.ContainerOp(\n",
    "        name = 'webapp',\n",
    "        image = 'gcr.io/google-samples/ml-pipeline-webapp-launcher',\n",
    "        arguments = [\"--model_name\", 'ghsumm-%s' % ('{{workflow.name}}',),\n",
    "            \"--github_token\", github_token]\n",
    "\n",
    "        )\n",
    "    webapp.after(serve)   \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile the new pipeline definition and submit the run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiler.Compiler().compile(gh_summ2, 'ghsumm2.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run = client.run_pipeline(exp.id, 'ghsumm2', 'ghsumm2.tar.gz',\n",
    "                          params={'working-dir': WORKING_DIR,\n",
    "                                  'github-token': GITHUB_TOKEN,\n",
    "                                  'deploy-webapp': DEPLOY_WEBAPP,\n",
    "                                  'project': PROJECT_NAME})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![The new pipeline.](https://storage.googleapis.com/amy-jo/images/datagen_t2t_pipeline.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When this new pipeline finishes running, you'll be able to see your generated processed data files in GCS under the path: `WORKING_DIR/<pipeline_name>/gh_data`. There isn't time in the workshop to pre-process the full dataset, but if there had been, we could have defined our pipeline to read from that generated directory for its training input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----------------------------\n",
    "Copyright 2018 Google LLC\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "     http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
